{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import itertools\n",
    "import pymc3 as pm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N,I,K,T = 3500,20,4,3\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q = list(itertools.product([0,1],repeat=4))\n",
    "Q = np.array(Q)[np.random.randint(1,2**4,I)]\n",
    "\n",
    "all_p = np.asarray(list(itertools.product([0,1],repeat=4)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "s = np.random.uniform(0,0.4,(1,I))\n",
    "g = np.random.uniform(0,0.4,(1,I))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta = np.random.normal(0,1,(N,1))\n",
    "lambda_k = np.random.lognormal(0,1,(1,I))\n",
    "lambda_0 = np.random.normal(0,1,(1,I))\n",
    "\n",
    "att_p_logit = theta*lambda_k+lambda_0\n",
    "att_p = np.exp(att_p_logit)/(1+np.exp(att_p_logit))\n",
    "att = np.random.binomial(1,att_p)\n",
    "eta = (att.reshape(N,1,K)>=Q.reshape(1,I,K)).prod(axis=-1)\n",
    "correct_p = eta*(1-s-g)+g\n",
    "Y = np.random.binomial(1,correct_p)\n",
    "c_eta = (all_p.reshape(-1,1,K)>=Q).prod(axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import theano.tensor as tt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q_tensor = tt.as_tensor(Q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pm.Model() as model:\n",
    "    theta_pre = pm.Normal(\"theta\",0,1,shape=(N,1))\n",
    "    lambda_0_pre = pm.Normal(\"lambda_0\",0,2,shape=(1,1))\n",
    "    lambda_k_pre = pm.Lognormal(\"lambda_k\",0,2,shape=(1,K))\n",
    "    att_p_pre = pm.math.invlogit(theta_pre*lambda_k_pre+lambda_0_pre)\n",
    "    c_pi = tt.exp(tt.dot(tt.log(att_p_pre),all_p.T)+tt.dot(tt.log(1-att_p_pre),1-all_p.T))\n",
    "    g_pre = pm.Beta(\"g\",1,1,shape=(1,I))\n",
    "    s_pre = pm.Uniform(\"s\",0,1-g_pre,shape=(1,I))\n",
    "    \n",
    "    pm.Deterministic(\"pi\",c_pi)\n",
    "    c_correct_prob = (1-s_pre-g_pre)*c_eta+g_pre\n",
    "    y_p = tt.dot(c_pi,c_correct_prob)\n",
    "    pm.Bernoulli(\"Y\",y_p,observed=Y)\n",
    "    # att_pre = pm.Bernoulli(\"att\",att_p_pre,shape=(N,K))\n",
    "    # eta = (att_pre>=Q_tensor).prod(axis=-1)\n",
    "\n",
    "    # Y_p = eta*(1-s-g)+g\n",
    "    # pm.Bernoulli(\"Y\",Y_p,observed=Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with model:\n",
    "    va = pm.ADVI(Minibatches=256)\n",
    "    approx = pm.fit(n=60000, method=va,obj_n_mc=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with model:\n",
    "    trace = pm.sample(2000,tune=2000,return_inferencedata=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import arviz as az"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr = approx.sample(draws=2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.abs(((az.summary(tr,\"s\")[\"mean\"].values-s))).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.abs((az.summary(tr,\"g\")[\"mean\"].values-g)).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.abs((az.summary(trace,\"s\")[\"mean\"].values-s)).mean())\n",
    "print(np.abs((az.summary(trace,\"g\")[\"mean\"].values-g)).mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.summary(trace,\"pi\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "att_p_pre.tag.test_value.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "att_p.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pi = np.exp(np.dot(np.log(att_p),all_p.T)+np.dot(np.log(1-att_p),1-all_p.T))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.dot(pi,c_eta).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_p[0]*t_v[0,0,]+(1-all_p[0])*(1-t_v[0,0,])"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "075496efbd0cd29f44e211505c4029a718f6409ae34c249d627c1cb29e30b628"
  },
  "kernelspec": {
   "display_name": "mcmc",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
